# components/storage.py

from llama_index.vector_stores.postgres import PGVectorStore
from llama_index.storage.docstore.mongodb import MongoDocumentStore
from llama_index.core import VectorStoreIndex
from ..config import settings

def get_pg_vector_store():
    """Initializes and returns the PGVectorStore."""
    print("🔌 [Storage] Connecting to PGVectorStore...")
    return PGVectorStore.from_params(
        database=settings.PG_DATABASE,
        host=settings.PG_HOST,
        password=settings.PG_PASSWORD,
        port=settings.PG_PORT,
        user=settings.PG_USER,
        table_name=settings.PG_TABLE_NAME,
        embed_dim=settings.PG_EMBED_DIM,
    )

def get_mongo_docstore():
    """Initializes and returns the MongoDocumentStore."""
    print("🔌 [Storage] Connecting to MongoDocumentStore...")
    return MongoDocumentStore.from_uri(
        uri=settings.MONGO_URI,
        db_name=settings.MONGO_DB_NAME
    )

def get_nodes_from_docstore(docstore: MongoDocumentStore):
    """Fetches all nodes from the document store."""
    print("🚚 [Storage] Fetching nodes from MongoDocumentStore...")
    return list(docstore.docs.values())

def get_vector_index(vector_store: PGVectorStore):
    """Builds a VectorStoreIndex from a vector store."""
    print("🏗️ [Storage] Building VectorStoreIndex from store...")
    return VectorStoreIndex.from_vector_store(vector_store=vector_store)